"""
检查 SV 是否在 block 中
"""

import os
import sys
import pandas as pd


SAMPLE_NAME = sys.argv[1]
DIR="/public/home/yunlzhang/data/pan-genome"

SV_DATA_FILE_PATH = "./SV_data.txt"
sv_df = pd.read_csv(SV_DATA_FILE_PATH)


def parse_block(chr_id, sv_start, sv_stop):
    """
    解析haploview的结果
    """
    sv_dir = f"{DIR}/{SAMPLE_NAME}/haploview/{chr_id}-{sv_start}"
    sv_id = f"{SAMPLE_NAME}-{chr_id}-{sv_start}"

    block_file = f"{sv_dir}/haploview.GABRIELblocks"
    info_file = f"{sv_dir}/SV.plink.info"
    all_snp_df = pd.DataFrame()
    all_block_context = []

    include_sv_block= {}

    if os.path.exists(block_file) and os.path.exists(info_file):
        info_df = pd.read_csv(info_file, header=None, sep="\t")
        block_count = 0
        with open(block_file, "r") as f:
            lines = f.readlines()
            block_context = []
            block_id = ""
            for line in lines:
                line_text = line.strip()
                if "BLOCK" in line_text:
                    block_count+=1
                    block_id = f"{sv_id}-{block_count}"
                    line_list = line_text.split(" ")
                    marker_index = line_list.index("MARKERS:")
                    snp_indexs = [int(v)-1 for v in line_list[marker_index+1:]]
                    item_snp_df = info_df.iloc[snp_indexs][[0]].copy()
                    item_snp_df[1] = block_id

                    # check if sv in block
                    block_range = [int(info_df.iloc[snp_indexs[0]][1]), int(info_df.iloc[snp_indexs[-1]][1])]
                    if block_range[0] <= sv_start and block_range[1] >= sv_stop:
                        include_sv_block = {
                            "block_id": block_id,
                            "start": block_range[0],
                            "stop": block_range[1],
                            "snps": len(snp_indexs)
                        }

                    if len(all_snp_df) > 0:
                        all_snp_df = pd.concat([all_snp_df, item_snp_df])
                    else:
                        all_snp_df = item_snp_df

                    if len(block_context) > 0:
                        all_block_context += block_context
                        block_context = []
                else:
                    if len(block_context) == 0 and block_id !="":
                        block_context.append(f">{block_id}")
                        block_context.append(line_text)
                    else:
                        block_context.append(line_text)
            if len(block_context) > 0:
                all_block_context += block_context

    return all_snp_df, all_block_context, "\t".join([str(v) for v in list(include_sv_block.values())])


def find_tag_snp(sv_dir, sv_snp_df):
    """
    找出block中的Tag SNP
    """
    tag_file = f"{sv_dir}/haploview.TAGS"
    snp_tag_df = sv_snp_df.set_index(0)
    snp_tag_df["Tag"] = "N"

    if os.path.exists(tag_file):
        skiprows = 0
        target_line_exists = False

        with open(tag_file, "r") as f:
            for line in f:
                if "Test	Alleles Captured" in line:
                    target_line_exists = True
                    break
                else:
                    skiprows+=1

        # check if target line exists
        if target_line_exists:
            tag_df = pd.read_csv(tag_file, sep="\t", skiprows=skiprows)
            for index, row in tag_df.iterrows():
                if row["Test"] in snp_tag_df.index:
                    if len(row["Alleles Captured"].split(",")) > 1:
                        snp_tag_df.at[row["Test"], 'Tag'] = "Y"

    snp_tag_df.reset_index(inplace=True)
    return snp_tag_df


if __name__ == "__main__":
    try:
        target_df = sv_df[(sv_df["check"] == True) & (sv_df["genome"] == SAMPLE_NAME)]

        output_filename = {
            "snp_file": f"{DIR}/{SAMPLE_NAME}/block_snp.txt",
            "block_file": f"{DIR}/{SAMPLE_NAME}/block_context.txt",
            "sv_block_file": f"{DIR}/{SAMPLE_NAME}/sv_block.txt"
        }

        for k, v in output_filename.items():
            if os.path.exists(v):
                os.remove(v)

        for index, row in target_df.iterrows():
            chr_id = row["#reference"]
            start = row["ref_start"]
            stop = row["ref_stop"]
            sv_dir = f"{DIR}/{SAMPLE_NAME}/haploview/{chr_id}-{start}"
            all_snp_df, all_block_context, include_sv_block = parse_block(chr_id, start, stop)
            if len(all_snp_df) > 0:
                # 找出Tag SNP
                tag_snp_df = find_tag_snp(sv_dir, all_snp_df)
                tag_snp_df.to_csv(output_filename["snp_file"], index=False, header=False, sep="\t", mode="a")

            if len(all_block_context) > 0:
                with open(output_filename["block_file"], "a+") as f:
                    for line in all_block_context:
                        f.write(f"{line}\n")

            if include_sv_block:
                with open(output_filename["sv_block_file"], "a+") as f:
                    f.write(f"{include_sv_block}\n")


    except Exception as e:
        print(e)
